import numpy as np
from matplotlib import pyplot as plt

Sigma_true = np.array([[2.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
                       [0.,  2., -1.,  0.,  0.,  0.,  0.,  0.],
                       [0., -1.,  2.,  0., -1.,  0.,  0., -1.],
                       [0.,  0.,  0.,  2.,  1.,  0.,  0.,  0.],
                       [0.,  0., -1.,  1.,  2.,  0.,  0.,  1.],
                       [0.,  0.,  0.,  0.,  0.,  2.,  0.,  0.],
                       [0.,  0.,  0.,  0.,  0.,  0.,  2.,  0.],
                       [0.,  0., -1.,  0.,  1.,  0.,  0.,  2.]])

X = np.array([[-0.3311197, -0.3932145,  0.7742677, -0.3097049, -1.5111846,
               0.3311429,  2.1538895,  0.883587],
              [-2.7057869, -0.2514156, -0.6271286,  0.546302,  1.6978879,
               -0.3421863, -0.6586413, -0.8402505],
              [-0.3192959,  1.2968608, -0.4446002,  1.2538625,  1.881935,
               -2.0727404, -1.997299,  0.3265276],
              [-0.4125173,  0.5194493, -1.4735353,  1.5864925,  0.2298139,
               0.8494314,  0.5313172, -0.7134051],
              [-2.7713921, -0.042382, -1.196607, -2.0107553, -0.2221868,
               -0.2953777, -1.7265336, -0.038322],
              [-1.0180134, -0.8166305,  0.4122348,  0.0026914, -0.9920446,
               2.0909459, -0.4258249, -1.0972184],
              [0.8650409,  1.3432579, -0.637971, -1.6605675, -2.1753728,
               0.9573123, -0.5445886,  0.4765308],
              [-0.2625614,  0.4310924, -1.1176806, -0.7204528,  1.6350725,
               0.6776547,  1.3796291,  0.522066],
              [-0.9123371, -1.9296817,  1.3247186,  1.4169735, -0.8366834,
               -0.5114306,  1.4192098, -0.6890962],
              [-0.4228602,  1.1161097,  1.6705814,  3.2483859,  1.6894376,
               -0.1231031,  1.1623457, -2.0525432],
              [-0.7095916, -1.2685852,  0.1599799, -0.4224032,  0.0252021,
               1.1433826, -0.7329448,  1.9290979],
              [-0.4633842, -0.0797081,  0.0926444,  0.2468722,  1.6129349,
               0.992853,  1.3698709,  0.4049546],
              [-0.5948823, -0.3180217,  0.6807469, -0.1785293,  0.2297931,
               2.0016365, -0.3317563, -0.4840257],
              [-0.1052824, -2.3678404, -0.6171589, -1.5935395,  1.1229458,
               -0.3642313,  0.2468903,  0.6588132],
              [-1.6527603, -2.1323678,  1.6783155,  3.0078456,  0.7259103,
               0.0490899,  0.4264523,  0.4042916],
              [3.0977721, -0.8027199,  0.2392123, -2.8330579, -0.912167,
               -0.8299413, -1.9825168, -0.1996475],
              [0.6697608,  1.0988143, -1.0885298,  0.8310039, -0.6695225,
               1.5023243,  0.0969627, -1.868516],
              [1.848579, -0.4324271, -1.7509966, -0.9744788,  1.2262814,
               -0.3216769, -1.740705,  0.5857962],
              [0.4199998,  0.1054119,  0.7350115,  1.7864235, -0.05206,
               -0.7381368, -1.8674077, -1.1380628],
              [0.6701009, -0.5945656, -1.3242998,  0.159206,  1.1155564,
               -2.6386036, -1.0102455,  0.2825183],
              [0.5837604, -1.4277136, -1.0824538,  0.9759325,  0.6190599,
               -1.3624196, -1.1606201,  2.3861869],
              [0.4824618,  1.4701, -0.9503183,  0.3830496, -0.5714352,
               0.1090383, -1.15373, -0.3257069],
              [-1.2122038, -1.6077674,  0.3301125,  1.4432326,  1.0597755,
               -0.8848272,  3.8469001, -0.8128089],
              [-2.1423175, -1.2859498, -0.3093076,  0.497371,  0.5315224,
               1.1975472, -0.1029956,  0.370612],
              [0.2174001,  1.2404896, -0.0527978, -1.70954, -0.5960344,
               1.2499632,  0.5449211, -0.0448455],
              [-1.3262849, -1.5899511,  0.7485016, -0.2622975, -0.4558745,
               1.9483187,  1.4903012,  0.8006743],
              [0.9212063, -4.7382379,  2.5271192,  3.109847,  1.5434335,
               -1.3491622,  1.6059323, -0.2300004],
              [-2.6407118,  1.5776585, -0.2052617,  0.3756392,  1.1194163,
               -3.2734215,  0.1157875,  0.7644046],
              [0.9610964, -1.1207784, -0.2378744,  1.0983257,  1.3869479,
               1.0115922, -1.5241613,  0.1350005],
              [0.2636862, -1.5156891, -0.8902854, -2.3923297, -0.5273719,
               2.8639841,  0.8964969,  1.3435503],
              [-0.4730562,  0.1234838, -0.9236915, -0.400059,  1.2714111,
               1.6974755,  1.2243627,  0.9603859],
              [1.4601264, -0.5659824,  1.9538219, -0.0001947, -1.106602,
               -3.0011821,  1.2980528, -3.6546994],
              [-0.0841816, -0.9661943,  0.2898518, -2.8143836, -0.3795513,
               1.3108776,  1.095498, -2.6934468],
              [0.1848951,  0.7326472, -0.0012979, -2.3559473, -0.4322298,
               0.6223169, -2.0225253,  2.3752747],
              [-0.8488349,  0.3068909,  1.5376214,  0.1497557, -1.2388319,
               -0.2815129,  0.653513, -1.9475487],
              [-2.760669, -2.1736976,  0.2587605, -0.6145264,  0.1003725,
               -3.0165647, -1.7505355,  1.0967079],
              [1.2015188,  0.8893968,  0.3391192, -0.6410825, -0.5730524,
               0.8334876, -0.6972086, -0.261479],
              [-0.7922156, -1.3573571,  1.2776654, -0.4787305, -1.3617914,
               1.172376,  0.8904296,  0.1550097]])

lam = .15
N = X.shape[0]
Sigma_emp = X.T@X/N


def plot_helper(Z, save_dest=None):
    cmap = "viridis"
    vmax = max([A.max() for A in [Z, Sigma_emp, Sigma_true]])
    vmin = max([A.min() for A in [Z, Sigma_emp, Sigma_true]])
    plt.subplots(1, 3, facecolor="w", figsize=(6, 6))
    plt.subplot(3, 1, 1)
    plt.imshow(Sigma_true, vmin=vmin, vmax=vmax, cmap=cmap)
    plt.title(r"$\Sigma^\mathrm{true}$")
    plt.subplot(3, 1, 2)
    plt.title(r"$\hat\Sigma$")
    plt.imshow(Z, vmin=vmin, vmax=vmax, cmap=cmap)
    plt.subplot(3, 1, 3)
    plt.title(r"$\Sigma^\mathrm{emp}$")
    plt.imshow(Sigma_emp, vmin=vmin, vmax=vmax, cmap=cmap)
    plt.tight_layout()

    if save_dest is not None:
        plt.savefig(save_dest, bbox_inches='tight')
    else:
        plt.show()